# -*- coding:utf-8 -*-
from rich import print  # 导入 rich 库，用于美化打印输出
from transformers import AutoTokenizer  # 导入 Hugging Face 的 AutoTokenizer，用于文本 tokenization
import numpy as np  # 导入 NumPy，用于处理数组和矩阵操作
import sys  # 导入 sys 模块，用于访问系统相关的功能（如路径等）

sys.path.append('..')  # 将上一级目录添加到模块搜索路径中，便于导入自定义模块
from pet_config import *  # 导入配置文件中的内容（如 ProjectConfig 类）


class HardTemplate(object):
    """
    硬模板类，用于人工定义句子和 [MASK] 之间的位置关系。
    """

    def __init__(self, prompt: str):
        """
        初始化方法，解析 Prompt 模板字符串。

        Args:
            prompt (str): 定义的 Prompt 格式字符串，例如 "这是一条{MASK}评论：{textA}。"
        """
        self.prompt = prompt  # 存储传入的 Prompt 字符串
        self.inputs_list = []  # 存储拆解后的 Prompt 各部分
        self.custom_tokens = set(['MASK'])  # 存储 Prompt 中涉及的自定义字段（如 {MASK} 和 {textA}）
        self.prompt_analysis()  # 调用方法解析 Prompt 模板

    def prompt_analysis(self):
        """
        解析 Prompt 模板字符串，将其拆解为可映射的数据结构。

        Examples:
            输入: "这是一条{MASK}评论：{textA}。"
            输出:
                inputs_list -> ['这', '是', '一', '条', 'MASK', '评', '论', '：', 'textA', '。']
                custom_tokens -> {'textA', 'MASK'}
        """
        idx = 0  # 初始化索引
        while idx < len(self.prompt):  # 遍历整个 Prompt 字符串
            str_part = ''  # 临时存储当前自定义字段的内容
            if self.prompt[idx] not in ['{', '}']:  # 如果当前字符不是 "{" 或 "}"，直接加入 inputs_list
                self.inputs_list.append(self.prompt[idx])
            if self.prompt[idx] == '{':  # 如果遇到 "{"，表示进入一个自定义字段
                idx += 1  # 移动到下一个字符
                while self.prompt[idx] != '}':  # 循环读取自定义字段的内容，直到遇到 "}"
                    str_part += self.prompt[idx]  # 拼接字段内容
                    idx += 1  # 继续向后移动
            elif self.prompt[idx] == '}':  # 如果单独遇到 "}"，说明括号不匹配
                raise ValueError("Unmatched bracket '}', check your prompt.")  # 抛出错误提示
            if str_part:  # 如果有拼接的字段内容
                self.inputs_list.append(str_part)  # 加入 inputs_list
                self.custom_tokens.add(str_part)  # 加入 custom_tokens 集合
            idx += 1  # 处理完当前字符后继续向后移动

    def __call__(self,
                 inputs_dict: dict,
                 tokenizer,
                 mask_length,
                 max_seq_len=512):
        """
        根据输入字典和 Tokenizer，将样本转换为符合模板格式的编码结果。

        Args:
            inputs_dict (dict): 包含 Prompt 中参数的字典，例如 {"textA": "包装不错", "MASK": "[MASK]"}
            tokenizer: 用于对文本进行编码的 Tokenizer 对象
            mask_length (int): MASK token 的重复次数
            max_seq_len (int): 最大序列长度，默认为 512

        Returns:
            dict: 包含以下键值的字典：
                - text: 编码后的原始文本
                - input_ids: Token ID 序列
                - token_type_ids: 分段 ID 序列
                - attention_mask: 注意力掩码序列
                - mask_position: MASK token 的位置索引列表
                Returns:
            dict -> {
                'text': '[CLS]这是一条[MASK]评论：这个手机也太卡了。[SEP]',
                'input_ids': [1, 47, 10, 7, 304, 3, 480, 279, 74, 47, 27, 247, 98, 105, 512, 777, 15, 12043, 2],
                'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                'mask_position': [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
            }
        """
        # 初始化输出字典
        outputs = {
            'text': '',
            'input_ids': [],
            'token_type_ids': [],
            'attention_mask': [],
            'mask_position': []
        }

        str_formated = ''  # 初始化格式化后的字符串
        for value in self.inputs_list:  # 遍历拆解后的 Prompt 各部分
            if value in self.custom_tokens:  # 如果是自定义字段
                if value == 'MASK':  # 如果是 MASK 字段
                    str_formated += inputs_dict[value] * mask_length  # 按指定长度重复 MASK
                else:  # 其他字段
                    str_formated += inputs_dict[value]  # 直接替换为对应的值
            else:  # 如果是普通字符
                str_formated += value  # 直接拼接到结果字符串中
        print(f'str_formated-->{str_formated}')  # 打印格式化后的字符串

        # 使用 Tokenizer 对格式化后的字符串进行编码
        encoded = tokenizer(
            text=str_formated,  # 输入文本
            truncation=True,  # 截断超出最大长度的部分
            max_length=max_seq_len,  # 最大序列长度
            padding='max_length'  # 填充至最大长度
        )
        print(f'encoded--->{encoded}')  # 打印编码结果

        # 更新输出字典
        outputs['input_ids'] = encoded['input_ids']  # Token ID 序列
        outputs['token_type_ids'] = encoded['token_type_ids']  # 分段 ID 序列
        outputs['attention_mask'] = encoded['attention_mask']  # 注意力掩码序列
        outputs['text'] = ''.join(tokenizer.convert_ids_to_tokens(encoded['input_ids']))  # 还原为文本
        mask_token_id = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]  # 获取 MASK token 的 ID
        mask_position = np.where(np.array(outputs['input_ids']) == mask_token_id)[0].tolist()  # 查找 MASK 的位置
        outputs['mask_position'] = mask_position  # 更新 MASK 位置列表
        return outputs  # 返回最终结果


if __name__ == '__main__':
    pc = ProjectConfig()  # 创建配置对象
    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)  # 加载预训练的 Tokenizer
    hard_template = HardTemplate(prompt='这是一条{MASK}评论：{textA}')  # 创建硬模板对象
    print(hard_template.inputs_list)  # 打印拆解后的 Prompt 列表
    print(hard_template.custom_tokens)  # 打印自定义字段集合
    tep = hard_template(  # 调用硬模板对象，生成编码结果
        inputs_dict={'textA': '包装不错，苹果挺甜的，个头也大。', 'MASK': '[MASK]'},  # 输入字典
        tokenizer=tokenizer,  # 使用的 Tokenizer
        max_seq_len=30,  # 最大序列长度
        mask_length=2  # MASK 长度
    )
    print(tep)  # 打印编码结果

    print(tokenizer.convert_ids_to_tokens([3819, 3352, 3819, 3352]))  # 测试 ID 转换为 Token
    print(tokenizer.convert_tokens_to_ids(['网', '球']))  # 测试 Token 转换为 ID